# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#:
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Calculate cross entropy loss"""
from mindspore import Tensor
import mindspore
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
import mindspore.ops.operations as P
from mindspore.nn.loss.loss import LossBase
from mindspore.ops import functional as F
import mindspore.numpy as np
from mindspore.nn.loss.loss import LossBase
from mindspore import ops
from mindspore.ops import functional as F
from mind3d.utils.class_factory import ClassFactory, ModuleType



__all__=["CrossEntropySmooth_SEG","CrossEntropySmooth_CLS",'NLLLoss']

class CrossEntropySmooth_CLS(LossBase):
    """CrossEntropy for cls"""
    def __init__(self, sparse=True, reduction='mean', smooth_factor=0., num_classes=1000):
        super(CrossEntropySmooth_CLS, self).__init__()
        self.onehot = ops.OneHot()
        self.sparse = sparse
        self.on_value = ms.Tensor(1.0 - smooth_factor, ms.float32)
        self.off_value = ms.Tensor(1.0 * smooth_factor / (num_classes - 1), ms.float32)
        self.ce = nn.SoftmaxCrossEntropyWithLogits(reduction=reduction)

    def construct(self, logit, label):
        if self.sparse:
            label = self.onehot(label, ops.shape(logit)[1], self.on_value, self.off_value)
        loss = self.ce(logit, label)
        return loss


class CrossEntropySmooth_SEG(LossBase):
    """CrossEntropy for seg"""
    def __init__(self, sparse=True, reduction='mean', smooth_factor=0., num_classes=1000):
        super(CrossEntropySmooth_SEG, self).__init__()
        self.onehot = ops.OneHot()
        self.sparse = sparse
        self.on_value = ms.Tensor(1.0 - smooth_factor, ms.float32)
        self.off_value = ms.Tensor(1.0 * smooth_factor / (num_classes - 1), ms.float32)
        self.ce = nn.SoftmaxCrossEntropyWithLogits(reduction=reduction)

    def construct(self, logit, label):

        logit=logit.transpose(0,2,1)
        logit=logit.view(-1,13)
        label=label.view(-1,1).squeeze()
        if self.sparse:
            label = self.onehot(label, ops.shape(logit)[1], self.on_value, self.off_value)
        loss = self.ce(logit, label)
        return loss

"""NLLLoss"""
@ClassFactory.register(ModuleType.LOSS)
class NLLLoss(LossBase):
    """NLLLoss"""

    def __init__(self, reduction="mean"):
        super(NLLLoss, self).__init__(reduction)
        self.one_hot = ops.OneHot()
        self.reduce_sum = ops.ReduceSum()
        self.reshape = ops.Reshape()

    def construct(self, logits, labels, weights=None, method=None):
        """NLLLoss construct."""

        if method is True:
            logits = self.reshape(logits, (-1, 50))
            labels = labels.view(-1, 1).squeeze()
            label_one_hot = self.one_hot(labels, F.shape(logits)[-1], F.scalar_to_array(1.0), F.scalar_to_array(0.0))
            loss = self.reduce_sum(-1.0 * logits * label_one_hot, (1,))
            return self.get_loss(loss)

        elif weights is not None:
            """WeightedNLLLoss construct."""
            logits = ops.LogSoftmax(1)(logits)
            label_one_hot = self.one_hot(labels, F.shape(logits)[-1], F.scalar_to_array(1.0), F.scalar_to_array(0.0))
            loss = self.reduce_sum(-1.0 * logits * label_one_hot, (1,))
            return self.get_loss(loss, weights)

        else:
            """Default NLLLoss construct."""
            label_one_hot = self.one_hot(labels, F.shape(logits)[-1], F.scalar_to_array(1.0), F.scalar_to_array(0.0))
            loss = self.reduce_sum(-1.0 * logits * label_one_hot, (1,))
            return self.get_loss(loss)
